package kvraft

import (
	"6.824/labgob"
	"6.824/labrpc"
	"6.824/raft"
	"bytes"
	"fmt"
	"log"
	"sync"
	"sync/atomic"
	"time"
)

const Debug = false
//const Debug = true

func DPrintf(format string, a ...interface{}) (n int, err error) {
	if Debug {
		log.Printf(format, a...)
	}
	return
}

type Op struct {
	// Your definitions here.
	// Field names must start with capital letters,
	// otherwise RPC will break.
	Type   string
	Key    string
	Val    string
	From   int
	Id     int
	Client int64
}

type KVServer struct {
	mu        sync.Mutex
	me        int
	rf        *raft.Raft
	persister *raft.Persister
	applyCh   chan raft.ApplyMsg
	dead      int32 // set by Kill()

	maxraftstate int // snapshot if log grows this big

	// Your definitions here.
	listen    int32
	data      map[string]string
	reqData   map[int64]record
	retCh     chan GetReply
	lastApply int
}

type record struct {
	Id  int
	Val string
}

type snapshot struct {
	Data    map[string]string
	ReqData map[int64]record
}

func (kv *KVServer) pushUnblock(data GetReply, ch chan GetReply) {
	if atomic.LoadInt32(&kv.listen) == 0 {
		return
	}
	select {
	case ch <- data:
		//case <-time.After(3 * time.Millisecond):

	}
}

func (kv *KVServer) commitWork(ch chan raft.ApplyMsg) {

	for !kv.killed() {
		msg := <-ch
		DPrintf("me[%v], msg[%v]", kv.me, msg)
		if !msg.CommandValid && !msg.SnapshotValid {
			kv.pushUnblock(GetReply{"leader change", kv.me, "", msg}, kv.retCh)
			continue
		}
		if msg.SnapshotValid && kv.rf.CondInstallSnapshot(msg.SnapshotTerm, msg.SnapshotIndex, msg.Snapshot) {
			r := bytes.NewBuffer(msg.Snapshot)
			d := labgob.NewDecoder(r)
			var xxx snapshot
			err := d.Decode(&xxx)
			if err != nil {
				DPrintf("snapshot err[%v]", err)
				continue
			}
			kv.data = xxx.Data
			kv.reqData = xxx.ReqData
			continue
		}
		if !msg.CommandValid {
			continue
		}
		command := msg.Command.(Op)
		val := command.Val
		key := command.Key

		//if msg.CommandIndex <= kv.lastApply {
		//	if command.From == kv.me {
		//		kv.retCh <- GetReply{Err: ErrLateLog}
		//	}
		//	return
		//}
		//kv.lastApply = msg.CommandIndex

		client := command.Client
		r, ok := kv.reqData[client]
		if !ok {
			r = record{0, ""}
		}

		if command.Id <= r.Id {

			//delete(kv.reqData, command.Id)

			kv.pushUnblock(GetReply{"", kv.me, r.Val, msg}, kv.retCh)
			continue
		}
		kv.lastApply = command.Id

		//kv.mu.Lock()

		if command.Type == PUT {
			kv.data[key] = val
		} else if command.Type == APPEND {
			s, ok := kv.data[key]
			if !ok {
				s = ""
			}
			kv.data[key] = s + val
		}
		val, ok = kv.data[key]
		if !ok {
			val = ""
		}
		kv.reqData[client] = record{command.Id, val}

		kv.pushUnblock(GetReply{Value: val, Msg: msg}, kv.retCh)

		if kv.maxraftstate != -1 && kv.maxraftstate <= kv.persister.RaftStateSize() {
			w := new(bytes.Buffer)
			e := labgob.NewEncoder(w)
			err := e.Encode(snapshot{Data: kv.data, ReqData: kv.reqData})
			if err != nil {
				DPrintf("kv snap err[%v]", err)
				continue
			}
			kv.rf.Snapshot(msg.CommandIndex, w.Bytes())
		}
	}
}

func (kv *KVServer) Get(args *GetArgs, reply *GetReply) {
	// Your code here.
	if kv.killed() {
		return
	}
	kv.mu.Lock()
	defer kv.mu.Unlock()

	op := Op{GET, args.Key, "", kv.me, args.Id, args.ClientId}
	reply.Server = kv.me
	now := time.Now()
	var ret GetReply
	defer func() {
		DPrintf("get cost[%v]", time.Now().Sub(now))
		DPrintf("get me[%v] arg[%v] ret[%v] reply[%v]", kv.me, args, ret, reply)
		atomic.StoreInt32(&kv.listen, 0)
	}()

	index, _, isLeader := kv.rf.Start(op)
	if !isLeader {
		reply.Err = "not leader"
		return
	}

	atomic.StoreInt32(&kv.listen, 1)

	//time.Sleep(1000 * time.Millisecond)
	//for time.Since(now) < 200*time.Millisecond {
	for {
		select {
		case ret = <-kv.retCh:
			if len(ret.Err) > 0 {
				reply.Err = ret.Err
				return
			}
			applyMsg := ret.Msg
			if applyMsg.CommandIndex < index {
				continue
			}
			if applyMsg.CommandIndex > index || applyMsg.Command != op {
				reply.Err = "loses its leadership"
			}
			reply.Value = ret.Value
			return
		}
	}

}

func (kv *KVServer) PutAppend(args *PutAppendArgs, reply *PutAppendReply) {
	// Your code here.
	if kv.killed() {
		return
	}

	kv.mu.Lock()
	defer kv.mu.Unlock()
	now := time.Now()
	op := Op{args.Op, args.Key, args.Value, kv.me, args.Id, args.ClientId}
	defer func() {
		DPrintf("[%v] me[%v] arg[%v],  reply[%v]", op.Type, kv.me, args, reply)
		DPrintf("PutAppend cost[%v]", time.Now().Sub(now))
		atomic.StoreInt32(&kv.listen, 0)

	}()

	reply.Server = kv.me
	index, term, isLeader := kv.rf.Start(op)
	if !isLeader {
		sprintf := fmt.Sprintf("me[%v] index[%v], term[%v], isLeader[%v]", kv.me, index, term, isLeader)
		reply.Err = Err(sprintf)
		return
	}

	atomic.StoreInt32(&kv.listen, 1)

	//for time.Since(now) < 200*time.Millisecond {
	for {
		select {
		case ret := <-kv.retCh:
			if len(ret.Err) > 0 {
				reply.Err = ret.Err
				return
			}
			applyMsg := ret.Msg
			if applyMsg.CommandIndex < index {
				continue
			}
			if applyMsg.CommandIndex > index || applyMsg.Command != op {
				reply.Err = "loses its leadership"
			}
			return
		}
		//DPrintf("put sleep")
	}

}

func serverpau() {
	time.Sleep(5 * time.Millisecond)
}

//
// the tester calls Kill() when a KVServer instance won't
// be needed again. for your convenience, we supply
// code to set rf.dead (without needing a lock),
// and a killed() method to test rf.dead in
// long-running loops. you can also add your own
// code to Kill(). you're not required to do anything
// about this, but it may be convenient (for example)
// to suppress debug output from a Kill()ed instance.
//
func (kv *KVServer) Kill() {
	atomic.StoreInt32(&kv.dead, 1)
	kv.rf.Kill()
	// Your code here, if desired.
}

func (kv *KVServer) killed() bool {
	z := atomic.LoadInt32(&kv.dead)
	return z == 1
}

func (kv *KVServer) readPersister(data []byte) {
	if data == nil || len(data) <= 0 {
		return
	}
	r := bytes.NewBuffer(data)
	d := labgob.NewDecoder(r)
	var xxx snapshot
	err := d.Decode(&xxx)
	if err != nil {
		DPrintf("snapshot err[%v]", err)
		return
	}
	kv.data = xxx.Data
	kv.reqData = xxx.ReqData
}

//
// servers[] contains the ports of the set of
// servers that will cooperate via Raft to
// form the fault-tolerant key/value service.
// me is the index of the current server in servers[].
// the k/v server should store snapshots through the underlying Raft
// implementation, which should call persister.SaveStateAndSnapshot() to
// atomically save the Raft state along with the snapshot.
// the k/v server should snapshot when Raft's saved state exceeds maxraftstate bytes,
// in order to allow Raft to garbage-collect its log. if maxraftstate is -1,
// you don't need to snapshot.
// StartKVServer() must return quickly, so it should start goroutines
// for any long-running work.
//
func StartKVServer(servers []*labrpc.ClientEnd, me int, persister *raft.Persister, maxraftstate int) *KVServer {
	// call labgob.Register on structures you want
	// Go's RPC library to marshall/unmarshall.
	labgob.Register(Op{})

	kv := new(KVServer)
	kv.me = me
	kv.maxraftstate = maxraftstate
	kv.persister = persister
	kv.data = make(map[string]string)
	kv.reqData = make(map[int64]record)
	kv.readPersister(persister.ReadSnapshot())
	// You may need initialization code here.

	kv.applyCh = make(chan raft.ApplyMsg)
	kv.retCh = make(chan GetReply)
	kv.rf = raft.Make(servers, me, persister, kv.applyCh)

	// You may need initialization code here.

	go kv.commitWork(kv.applyCh)
	return kv
}
